import sys
import getopt
import numpy as np
import mindspore.context as context
from mindspore import Tensor
from mindspore import nn 
import mindspore.ops.operations as P
import mindspore.ops.functional as F

opts,_ = getopt.getopt(sys.argv[1:], "d:t:")
opts = dict(opts)
dev_id = 0 if '-d' not in opts else int(opts['-d'])
test_type = 0 if '-t' not in opts else int(opts['-t'])

context.set_context(mode=context.GRAPH_MODE, device_target="GPU", device_id=dev_id)
if test_type > 0:
    context.set_context(enable_graph_kernel=True)
if test_type > 1:
    context.set_context(graph_kernel_flags="--enable_stitch_fusion")


class EmbeddingPostprocessor(nn.Cell):
    def __init__(self):
        super(EmbeddingPostprocessor, self).__init__()
        self.layernorm = nn.LayerNorm((768,))
        self.add = P.Add()
        self.dropout = nn.Dropout(1-0.1)

    def construct(self, word_embeddings, token_type_embeddings, position_embeddings):
        output = word_embeddings
        output = self.add(output, token_type_embeddings)
        output = self.add(output, position_embeddings)
        output = self.layernorm(output)
        output = self.dropout(output)
        return output

shape1 = [8192, 768]
shape2 = [1, 768]
x = Tensor(np.random.normal(0, 1, shape1).astype(np.float32))
y = Tensor(np.random.normal(0, 1, shape1).astype(np.float32))
z = Tensor(np.random.normal(0, 1, shape2).astype(np.float32))
net = EmbeddingPostprocessor()
result = net(x,y,z)
